# encoding: utf-8

import os

from torch import optim

from constant import *
from models import *

# 初始化UNet-256。如果是UNet-128，则下采样次数为7
# facade数据集，输入输出都是RGB图像
netG = UnetGenerator(input_nc, output_nc, 8, ngf).to(device)

# 使用多个gpu处理
if (device.type == 'cuda') and (ngpu > 1):
    netG = nn.DataParallel(netG, list(range(ngpu)))

# 使用自定义的权重初始化函数
netG.apply(weights_init)

# 打印模型
# print(netG)

netD = NLayerDiscriminator(input_nc+output_nc, ndf, 3).to(device)

# 使用多个gpu处理
if (device.type == 'cuda') and (ngpu > 1):
    netD = nn.DataParallel(netD, list(range(ngpu)))

# 应用weights_init函数随机初始化所有权重，mean= 0，stdev = 0.2
netD.apply(weights_init)

# 打印模型
# print(netD)


# 初始化BCELoss函数
criterion = nn.BCEWithLogitsLoss()

# 生成器的L1正则化项
criterionL1 = torch.nn.L1Loss()


# 定义训练期间真假图像的标签值
real_label = torch.tensor(1, dtype=torch.float, device=device)
fake_label = torch.tensor(0, dtype=torch.float, device=device)

# 为 G 和 D 创建 Adam 优化器
optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))


def set_requires_grad(nets, requires_grad=False):
    """将网络的所参数设置 requies_grad=Fasle以减少不必要的计算
    参数:
        nets (network list)   -- 网络列表
        requires_grad (bool)  -- 是否需要计算梯度
    """
    if not isinstance(nets, list):
        nets = [nets]
    for net in nets:
        if net is not None:
            for param in net.parameters():
                param.requires_grad = requires_grad


def make_label(prediction, target_is_real):
    if target_is_real:
        target_tensor = real_label
    else:
        target_tensor = fake_label
    return target_tensor.expand_as(prediction)


# 记录训练的过程
img_list = []
G_losses = []
D_losses = []
iters = 0

# 取第一张图为固定的约束，以观察生成器的学习过程
fixed_A, fixed_B = None, None

netG.train()
netD.train()


def save_weights(epoch, netG, errG_GAN, errG_L1):
    if not os.path.exists("./weights"):
        os.mkdir("./weights")
    file_path = f"./weights/netG_epoch{epoch}_err{errG_GAN:.5f}_l1{errG_L1:0.5f}.pth"
    torch.save(netG.state_dict(), file_path)

"""
[  5/200] errD_fake: 0.5765  errD_real: 0.3845  errG_GAN: 1.3004  errG_L1: 66.7705
[  6/200] errD_fake: 0.4850  errD_real: 0.3999  errG_GAN: 1.4588  errG_L1: 63.6979
[  7/200] errD_fake: 0.3992  errD_real: 0.3316  errG_GAN: 1.5692  errG_L1: 60.8018
[  8/200] errD_fake: 0.3528  errD_real: 0.3176  errG_GAN: 1.6967  errG_L1: 58.0447
"""


print("Starting Training Loop...")
# 遍历数据集
for epoch in range(num_epochs):
    # 遍历真实数据中的每一个数据批
    for i, (real_A, real_B) in enumerate(dataloader, 0):

        real_A = real_A.to(device)
        real_B = real_B.to(device)

        if fixed_A is None:
            fixed_A = real_A[:1, :, :, :]
            fixed_B = real_B[:1, :, :, :]

        ############################
        # (1) 更新判别器: 最大化 log(D(x)) + log(1 - D(G(z)))
        ###########################

        set_requires_grad(netD, True)  # 需要计算D的梯度
        netD.zero_grad()  # 清空D的梯度

        ## 使用真实数据构建训练批
        # 拼接图像和约束
        real_AB = torch.cat((real_A, real_B), 1)
        # 预测真实数据
        pred_real = netD(real_AB)
        # 制作用于计算损失的label
        label = make_label(pred_real, True)
        # 计算真实数据的损失
        errD_real = criterion(pred_real, label)

        ## 使用生成器产生的数据构建训练批
        # 使用生成器产生假数据
        fake_B = netG(real_A)
        # 拼接 假图和约束
        fake_AB = torch.cat((real_A, fake_B), 1)
        # 通过将fake_B从计算图中剥离来隔断梯度传播到G
        pred_fake = netD(fake_AB.detach())
        # 制作用于计算损失的label
        label = make_label(pred_fake, False)
        # 计算假数据的损失
        errD_fake = criterion(pred_fake, label)

        # 合并损失，计算梯度
        errD = (errD_real + errD_fake) * 0.5
        errD.backward()

        # 更新判别器
        optimizerD.step()

        ############################
        # (2) 更新生成器: 最大化 log(D(G(z)))
        ###########################
        set_requires_grad(netD, False)  # 不需要 计算D的梯度
        netG.zero_grad()  # 清空生成器的梯度

        # 首先生成假图
        fake_AB = torch.cat((real_A, fake_B), 1)
        pred_fake = netD(fake_AB)
        # 制作用于计算损失的label，并计算损失
        label = make_label(pred_fake, True)  # 假数据在生成器看来是真数据
        errG_GAN = criterion(pred_fake, label)

        # 合并损失，计算梯度
        errG_L1 = criterionL1(fake_B, real_B) * lambda_L1
        errG = errG_GAN + errG_L1

        errG.backward()
        # 更新生成器
        optimizerG.step()

        # 保存训练时的损失，用于后续打印图形
        G_losses.append(errG.item())
        D_losses.append(errD.item())

        # 记录生成器将fixed_noise映射成了什么样的图片
        if (iters % 100 == 0) or ((epoch == num_epochs - 1) and (i == len(dataloader) - 1)):
            with torch.no_grad():
                fake = netG(fixed_A).detach().cpu()
            #             img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
            img_list.append(fake)

        iters += 1

    # 输出训练状态
    print('[%3d/%3d] errD_fake: %.4f  errD_real: %.4f  errG_GAN: %.4f  errG_L1: %.4f'
          % (epoch + 1, num_epochs,
             errD_fake.item(), errD_real.item(), errG_GAN.item(), errG_L1.item()))

    save_weights(epoch, netG, errG_GAN.item(), errG_L1.item())
